#include "gtest/gtest.h"
//#include "DataStructure/Common/include/Constant.h"
#include "sgArpack.h"

//using SG::DataStructure::Common::TOLERANCE;
using namespace SG::Algebra;

class ArpackTest : public ::testing::Test
{
  public:
    static void SetUpTestSuite ()
    {}
};

/* test program to solve for the 9 smallest eigenvalues of
 * A*x = lambda*x where A is the diagonal matrix
 * with entries 100, 99, ... , 2, 1 on the diagonal.
 */
TEST_F (ArpackTest, eig3_d_lm)
{
    typedef Eigen::Triplet<Real_t> Triplet;

    const int N   = 100;
    const int nev = 9;

    SpMat_t A (N, N);
    {
        std::vector<Triplet> triplets;
        for (int i = 0; i < N; ++i)
        {
            triplets.push_back (Triplet (i, i, N - i));
        }
        A.setFromTriplets (triplets.begin (), triplets.end ());
    }

    SpMat_t B (N, N);
    {
        std::vector<Triplet> triplets;
        for (int i = 0; i < N; ++i)
        {
            triplets.push_back (Triplet (i, i, 1));
        }
        B.setFromTriplets (triplets.begin (), triplets.end ());
    }

    // solve the generalized eigenvalue problem A*x = lambda*B*x
    TArpack<Real_t> solver;
    auto            opt = solver.getDefaultOption ();
    // SM for OP= (inv[K - sigma*M])*M, while LA for K*x = lambda*M*x
    opt.which = "LM";
    opt.nev   = 9;
    solver.setOption (opt);

    std::vector<Real_t> eigenvalues;
    std::vector<Vec_t>  eigenvectors;
    solver.solve (A, B, eigenvalues, eigenvectors);

    // Check
    auto info = solver.info ();
    ASSERT_TRUE (Info_t::Success == info);

    // the exact eigenvalues
    std::vector<Real_t> exact_eigenvalues;
    for (int i = 0; i < nev; ++i)
    {
        exact_eigenvalues.push_back (i + 1);
    }

    // // check the eigenvalues
    // // The converged Ritz values are always returned in ascending algebraic order
    // // see https://github.com/opencollab/arpack-ng/blob/master/SRC/dsaupd.f
    for (int i = 0; i < nev; ++i)
    {
        ASSERT_NEAR (exact_eigenvalues[i], eigenvalues[i], TOLERANCE);
    }
}

/* test program to solve for the 9 largest eigenvalues of
 * A*x = lambda*x where A is the diagonal matrix
 * with entries 100, 99, ... , 2, 1 on the diagonal.
 */
TEST_F (ArpackTest, eig3_d_sm)
{
    typedef Eigen::Triplet<Real_t> Triplet;

    const int N   = 100;
    const int nev = 9;

    SpMat_t A (N, N);
    {
        std::vector<Triplet> triplets;
        for (int i = 0; i < N; ++i)
        {
            triplets.push_back (Triplet (i, i, N - i));
        }
        A.setFromTriplets (triplets.begin (), triplets.end ());
    }

    SpMat_t B (N, N);
    {
        std::vector<Triplet> triplets;
        for (int i = 0; i < N; ++i)
        {
            triplets.push_back (Triplet (i, i, 1));
        }
        B.setFromTriplets (triplets.begin (), triplets.end ());
    }

    // solve the generalized eigenvalue problem A*x = lambda*B*x
    TArpack<Real_t> solver;
    auto            opt = solver.getDefaultOption ();
    // SM for OP= (inv[K - sigma*M])*M, while LA for K*x = lambda*M*x
    opt.which = "SM";
    opt.nev   = 9;
    solver.setOption (opt);

    std::vector<Real_t> eigenvalues;
    std::vector<Vec_t>  eigenvectors;
    solver.solve (A, B, eigenvalues, eigenvectors);

    // Check
    auto info = solver.info ();
    ASSERT_TRUE (Info_t::Success == info);

    // the exact eigenvalues
    std::vector<Real_t> exact_eigenvalues;
    for (int i = 0; i < nev; ++i)
    {
        exact_eigenvalues.push_back (N - nev + i + 1);
    }

    // // check the eigenvalues
    // // The converged Ritz values are always returned in ascending algebraic order
    // // see https://github.com/opencollab/arpack-ng/blob/master/SRC/dsaupd.f
    for (int i = 0; i < nev; ++i)
    {
        ASSERT_NEAR (exact_eigenvalues[i], eigenvalues[i], TOLERANCE);
    }
}

/* test program to solve for the 9 smallest eigenvalues of
 * A*x = lambda*x where A is the diagonal matrix
 * with entries 100, 99, ... , 2, 1 on the diagonal.
 *
 */
TEST_F (ArpackTest, eig3_z_lm)
{
    typedef Eigen::Triplet<Complex_t> Triplet;

    const int N   = 100;
    const int nev = 9;

    ComplexSpMat_t A (N, N);
    {
        std::vector<Triplet> triplets;
        for (int i = 0; i < N; ++i)
        {
            Complex_t value(N - i, 0.0);
            triplets.push_back (Triplet (i, i, value));
        }
        A.setFromTriplets (triplets.begin (), triplets.end ());
    }

    ComplexSpMat_t B (N, N);
    {
        std::vector<Triplet> triplets;
        for (int i = 0; i < N; ++i)
        {
            Complex_t value(1.0, 0.0);
            triplets.push_back (Triplet (i, i, value));
        }
        B.setFromTriplets (triplets.begin (), triplets.end ());
    }

    // solve the generalized eigenvalue problem A*x = lambda*B*x
    TArpack<Complex_t> solver;
    auto                               opt = solver.getDefaultOption ();
    // SM for OP= (inv[K - sigma*M])*M, while LA for K*x = lambda*M*x
    opt.which = "LM";
    opt.nev   = 9;
    solver.setOption (opt);

    std::vector<Complex_t>    eigenvalues;
    std::vector<ComplexVec_t> eigenvectors;

    solver.solve (A, B, eigenvalues, eigenvectors);
    

    // Check
    auto info = solver.info ();
    ASSERT_TRUE (Info_t::Success == info);

    // the exact eigenvalues
    std::vector<Complex_t> exact_eigenvalues;
    for (int i = 0; i < nev; ++i)
    {
        Complex_t value(i+1, 0.0);
        exact_eigenvalues.push_back (i + 1);
    }

    // // check the eigenvalues
    // // For TArpack<double>, it seems that the converged Ritz values are returned in
    // // a way which is different from TArpack<double>.
    // // see https://github.com/opencollab/arpack-ng/blob/master/SRC/znaupd.f
    for (int i = 0; i < nev; ++i)
    {
        ASSERT_NEAR (exact_eigenvalues[i].real (), eigenvalues[i].real (), TOLERANCE);
        ASSERT_NEAR (exact_eigenvalues[i].imag (), eigenvalues[i].imag (), TOLERANCE);
    }
}

/* test program to solve for the 9 largest eigenvalues of
 * A*x = lambda*x where A is the diagonal matrix
 * with entries 100, 99, ... , 2, 1 on the diagonal.
 */
TEST_F (ArpackTest, eig3_z_sm)
{
    typedef Eigen::Triplet<Complex_t> Triplet;

    const int N   = 100;
    const int nev = 9;

    ComplexSpMat_t A (N, N);
    {
        std::vector<Triplet> triplets;
        for (int i = 0; i < N; ++i)
        {
            Complex_t value (N - i, 0.0);
            triplets.push_back (Triplet (i, i, value));
        }
        A.setFromTriplets (triplets.begin (), triplets.end ());
    }

    ComplexSpMat_t B (N, N);
    {
        std::vector<Triplet> triplets;
        for (int i = 0; i < N; ++i)
        {
            Complex_t value (1.0, 0.0);
            triplets.push_back (Triplet (i, i, value));
        }
        B.setFromTriplets (triplets.begin (), triplets.end ());
    }

    // solve the generalized eigenvalue problem A*x = lambda*B*x
    TArpack<Complex_t> solver;
    auto                                 opt = solver.getDefaultOption ();
    // SM for OP= (inv[K - sigma*M])*M, while LA for K*x = lambda*M*x
    opt.which = "SM";
    opt.nev   = 9;
    solver.setOption (opt);

    std::vector<Complex_t>    eigenvalues;
    std::vector<ComplexVec_t> eigenvectors;

    solver.solve (A, B, eigenvalues, eigenvectors);

    // Check
    auto info = solver.info ();
    ASSERT_TRUE (Info_t::Success == info);

    // the exact eigenvalues
    std::vector<Complex_t> exact_eigenvalues;
    for (int i = 0; i < nev; ++i)
    {
        Complex_t value (N - i);
        exact_eigenvalues.push_back (value);
    }

    // check the eigenvalues
    // For TArpack<Complex_t>, it seems that the converged Ritz values are returned in
    // a way which is different from TArpack<double>.
    // see https://github.com/opencollab/arpack-ng/blob/master/SRC/znaupd.f
    for (int i = 0; i < nev; ++i)
    {
        ASSERT_NEAR (exact_eigenvalues[i].real (), eigenvalues[i].real (), TOLERANCE);
        ASSERT_NEAR (exact_eigenvalues[i].imag (), eigenvalues[i].imag (), TOLERANCE);
    }
}

// /* test program to solve for the 9 smallest eigenvalues of
//  * A*x = lambda*x where A is the diagonal matrix
//  * with entries 100, 99, ... , 2, 1 on the diagonal.
//  */
// TEST_F (ArpackTest, eig4_d_sm)
// {
//     typedef Eigen::Triplet<Real_t> Triplet;

//     const int N   = 100;
//     const int nev = 9;

//     SpMat_t A (N, N);
//     {
//         std::vector<Triplet> triplets;
//         for (int i = 0; i < N; ++i)
//         {
//             triplets.push_back (Triplet (i, i, N - i));
//         }
//         A.setFromTriplets (triplets.begin (), triplets.end ());
//     }

//     SpMat_t B (N, N);
//     {
//         std::vector<Triplet> triplets;
//         for (int i = 0; i < N; ++i)
//         {
//             triplets.push_back (Triplet (i, i, 1));
//         }
//         B.setFromTriplets (triplets.begin (), triplets.end ());
//     }

//     // solve the generalized eigenvalue problem A*x = lambda*B*x
//     TArpack<Real_t> solver;
//     auto            opt = solver.getDefaultOption ();
//     // SM for OP= (inv[K - sigma*M])*M, while LA for K*x = lambda*M*x
//     opt.mode = arpack::Mode::MODE_4;
//     opt.which = "SM";
//     opt.nev   = 9;
//     solver.setOption (opt);

//     std::vector<Real_t> eigenvalues;
//     std::vector<Vec_t>  eigenvectors;
//     solver.solve (A, B, eigenvalues, eigenvectors);

//     // Check
//     auto info = solver.info ();
//     ASSERT_TRUE (Info_t::Success == info);

//     // the exact eigenvalues
//     std::vector<Real_t> exact_eigenvalues;
//     for (int i = 0; i < nev; ++i)
//     {
//         exact_eigenvalues.push_back (i + 1);
//     }

//     // // check the eigenvalues
//     // // The converged Ritz values are always returned in ascending algebraic order
//     // // see https://github.com/opencollab/arpack-ng/blob/master/SRC/dsaupd.f
//     for (int i = 0; i < nev; ++i)
//     {
//         ASSERT_NEAR (exact_eigenvalues[i], eigenvalues[i], TOLERANCE);
//     }
// }